package medium.number;

/**
 * <a href="https://leetcode.cn/problems/find-the-punishment-number-of-an-integer/description/">2698. 求一个整数的惩罚数</a>
 * 给你一个正整数 n ，请你返回 n 的 惩罚数 。
 * n 的 惩罚数 定义为所有满足以下条件 i 的数的平方和：
 *   1 <= i <= n
 *   i * i 的十进制表示的字符串可以分割成若干连续子字符串，且这些子字符串对应的整数值之和等于 i 。
 * 示例 1：
 *   输入：n = 10
 *   输出：182
 *   解释：总共有 3 个整数 i 满足要求：
 *     - 1 ，因为 1 * 1 = 1
 *     - 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。
 *     - 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。
 *   因此，10 的惩罚数为 1 + 81 + 100 = 182
 * 示例 2：
 *   输入：n = 37
 *   输出：1478
 *   解释：总共有 4 个整数 i 满足要求：
 *     - 1 ，因为 1 * 1 = 1
 *     - 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。
 *     - 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。
 *     - 36 ，因为 36 * 36 = 1296 ，且 1296 可以分割成 1 + 29 + 6 。
 *   因此，37 的惩罚数为 1 + 81 + 100 + 1296 = 1478
 * 提示：
 *   1 <= n <= 1000
 * @author 刘学松
 * @date 2023-10-25 11:13
 */
public class 求一个整数的惩罚数 {
    public int punishmentNumber(int n) {
        // TODO 记忆优化
        int sum = n > 9 ? 182 : n == 9 ? 82 : 1;
        for (int i = 36; i <= n; i++) {
            int num = i * i;
            if (check(num, i, getMax(i))) {
                sum += num;
            }
        }
        return sum;
    }

    public boolean check(int num, int target, int max) {
        if (target < 0) {
            return false;
        }
        if (num == target) {
            return true;
        }
        for (int i = max; i > 1; i/=10) {
            int rightPart = num % i;
            int leftPart = num / i;
            int newTarget = target - leftPart;
            if (newTarget < 0) {
                break;
            }
            if (check(rightPart, newTarget, i/10)) {
                return true;
            }
        }
        return false;
    }

    public int getMax(int num) {
        if (num < 10000) {
            return 1000;
        } else if (num < 100000) {
            return 10000;
        } else if (num < 1000000) {
            return 100000;
        }
        return 1000000;
    }
}
